import sensor, time, ml, uos, gc
from machine import UART  # 新增：导入UART模块

sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QVGA)
sensor.set_windowing((240, 240))
sensor.skip_frames(time=2000)

# 新增：初始化串口（根据设备型号选择，此处以UART3为例，波特率115200）
uart = UART(3, 9600)
# uart = UART(1, 9600)  # 若为OpenMV RT，启用此句替换上一行

net = None
labels = None

try:
    net = ml.Model("trained.tflite", load_to_fb=uos.stat('trained.tflite')[6] > (gc.mem_free() - (64*1024)))
except Exception as e:
    print(e)
    raise Exception('Failed to load "trained.tflite" (' + str(e) + ')')

try:
    labels = [line.rstrip('\n') for line in open("labels.txt")]
except Exception as e:
    raise Exception('Failed to load "labels.txt" (' + str(e) + ')')

clock = time.clock()
while(True):
    clock.tick()
    img = sensor.snapshot()
    predictions_list = list(zip(labels, net.predict([img])[0].flatten().tolist()))

    # 新增：找到概率最高的类别
    max_pred = max(predictions_list, key=lambda x: x[1])  # 按概率排序取最大

    # 新增：根据最高概率类别发送对应数字（对应bottle、milk、napkin）
    if max_pred[0] == 'bottle':
        uart.write(bytes([1]))  # 发送数值1
        print(1)
    elif max_pred[0] == 'milk':
        uart.write(bytes([2]))  # 发送数值2
        print(2)
    elif max_pred[0] == 'napkin':
        uart.write(bytes([3]))  # 发送数值3
        print(3)

    # 保留原有打印功能
    for i in range(len(predictions_list)):
        print("%s = %f" % (predictions_list[i][0], predictions_list[i][1]))
    print(clock.fps(), "fps")
